#!/usr/bin/env python
# -*- coding: utf-8 -*-
# title       :
# author      : jarry
import os
# from unipath import Path
# from jarryLibs.fileLib.JPath import JPath
from JTK.jcore.JPath import JPath
import maya.cmds as cmds
import arnold
import mtoa

from . import __file__ as _initpy_file
__CURRENT_PATH__ = os.path.dirname(os.path.realpath(_initpy_file))


def getDefaultAllAOVs():
    builtin_aovs = [a[0] for a in mtoa.aovs.BUILTIN_AOVS]
    lighting_aovs = mtoa.aovs.LIGHTING_AOVS

    default_all_aovs = builtin_aovs + lighting_aovs
    return default_all_aovs


def createAOVs(aovs, driver_type, prefix='', half=True, merge=True,  lpe_list=None):
    cmds.lockNode('initialShadingGroup', l=0, lockUnpublished=0)
    cmds.lockNode('defaultTextureList1', l=0, lockUnpublished=0)
    cmds.lockNode('initialParticleSE', l=0, lockUnpublished=0)
    if cmds.getAttr("defaultRenderGlobals.currentRenderer") != 'arnold':
        cmds.setAttr("defaultRenderGlobals.currentRenderer", "arnold", type="string")

    aov_interface = mtoa.aovs.AOVInterface()
    for aovName in aovs:
        if aovName in ['crypto_asset', 'crypto_material', 'crypto_object']:
            create_cryptomatte_shader()
            cmds.setAttr('AOV_cryptomatte.userCryptoAov0', "crypto_abc", type="string")
            cmds.setAttr('AOV_cryptomatte.userCryptoSrc0', "abc", type="string")
            cmds.setAttr('AOV_cryptomatte.userCryptoAov1', "crypto_name", type="string")
            cmds.setAttr('AOV_cryptomatte.userCryptoSrc1', "name", type="string")
        aov_node_name = 'aiAOV_' + aovName
        if not cmds.objExists(aov_node_name):
            aov_res = aov_interface.addAOV(aovName)
            aov_node = aov_res.node

        else:
            aov_node = aov_node_name
        out = '{}.outputs[0]'.format(aov_node)

        # shader_name = ''
        if aovName == 'UV':
            shader_name = create_UV_shader()
            cmds.connectAttr("%s.outColor" % shader_name, "%s.defaultValue" % aov_node, f=1)
        else:
            # print aovName
            shader_name = 'AOV_%s_Shader' % aovName
            shader_file = JPath(__CURRENT_PATH__).child('customAovShader').child(shader_name + '.ma')
            if shader_file.exists() and not cmds.objExists(shader_name):
                # cmds.file(shader_file, i=1)
                cmds.file(shader_file, ignoreVersion=True, ra=True, mergeNamespacesOnClash=True, namespace=":",
                          options="v=0;p=17;f=0", pr=True, i=1)
            if cmds.objExists(shader_name):
                cmds.connectAttr("%s.outColor" % shader_name, "%s.defaultValue" % aov_node, f=1)

        aov_filter = cmds.listConnections('%s' % aov_node, d=0, type='aiAOVFilter')
        if aov_filter:
            aov_filter = aov_filter[0]
            if aov_filter != 'defaultArnoldFilter':
                cmds.delete(aov_filter)
                filterAttr = 'defaultArnoldFilter.message'
                cmds.connectAttr(filterAttr, '{}.filter'.format(out), f=1)

        aov_driver = cmds.listConnections('%s' % aov_node, d=0, type='aiAOVDriver')
        if aov_driver:
            aov_driver = aov_driver[0]
        # print driver_type
        if driver_type != 'Default':
            if cmds.getAttr('defaultArnoldDriver.prefix') != prefix:
                driver_node_name = 'aiAOVDriver_%s' % driver_type
                if cmds.objExists(driver_node_name):
                    driver_node = driver_node_name
                else:
                    driver_node = cmds.createNode('aiAOVDriver', name=driver_node_name, skipSelect=True)
                cmds.connectAttr('%s.message' % driver_node, '{}.driver'.format(out), f=1)
            else:
                driver_node = 'defaultArnoldDriver'
        else:
            driver_node = 'defaultArnoldDriver'
        # if prefix:
            # cmds.setAttr('%s.prefix' % driver_node, prefix.replace('<Group>', driver_type), type='string')
        cmds.setAttr('%s.prefix' % driver_node, prefix, type='string')
        cmds.setAttr('%s.halfPrecision' % driver_node, half)
        cmds.setAttr('%s.mergeAOVs' % driver_node, merge)
        lpe = ''
        if lpe_list:
            for l in lpe_list:
                if aovName in l:
                    lpe = l[aovName]
        cmds.setAttr('%s.lightPathExpression' % aov_node, lpe, type='string')

def removeAOVs(aovs):
    aov_interface = mtoa.aovs.AOVInterface()
    for aov in aovs:
        aov_interface.removeAOV(aov)


def create_cryptomatte_shader():
    shader_name = 'AOV_cryptomatte'
    if not cmds.objExists(shader_name):
        cmds.shadingNode('cryptomatte', au=True, name=shader_name)
        connect_list = cmds.listConnections('defaultArnoldRenderOptions.aovShaders', d=1)
        connect_number = len(connect_list) if connect_list else 0
        cmds.connectAttr('%s.outColor' % shader_name, 'defaultArnoldRenderOptions.aovShaders[%d]' % connect_number)
        return shader_name
    else:
        return shader_name


def create_UV_shader():
    shader_name = 'AOV_UV_Shader'
    if not cmds.objExists(shader_name):
        cmds.shadingNode('aiUtility', au=True, name=shader_name)
        cmds.setAttr('%s.shadeMode' % shader_name, 2)
        cmds.setAttr('%s.colorMode' % shader_name, 5)
    return shader_name


if __name__ == '__main__':
    # from ArnoldAOVsMG.data import utils as aovut
    #
    # reload(aovut)
    # pre = '<Scene>/<RenderLayer>/<RenderLayer>_{type}/<RenderLayer>_{type}'
    #
    # color_aovs = ['crypto_asset', 'crypto_material', 'crypto_object']
    # driver_type = 'cryptoMatte'
    # aovut.createAOVs(color_aovs, driver_type, pre)

    data_aovs = ['N', 'P', 'Z', 'opacity']
    driver_type = 'Data'
    pre = '<Scene>/<RenderLayer>/<RenderLayer>_{type}/<RenderLayer>_{type}'
    createAOVs(data_aovs, driver_type, pre)

    color_aovs = ['albedo', 'diffuse', 'diffuse_direct', 'diffuse_indirect', 'emission', 'specular', 'specular_direct',
                  'specular_indirect', 'sss', 'transmission', 'coat', 'direct']
    driver_type = 'RGBA'
    createAOVs(color_aovs, driver_type, pre)

    color_aovs = ['crypto_asset', 'crypto_material', 'crypto_object']
    driver_type = 'cryptoMatte'
    createAOVs(color_aovs, driver_type, pre)
